#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
小样本非参数统计MCP服务器 (n <= 30)
提供基于精确分布的单样本位置检验
完全独立实现，不依赖外部统计模块
"""

import asyncio
import json
import logging
import sys
from typing import Any, Dict, List, Tuple
import traceback
from math import factorial, comb, sqrt, pi, exp, log
from itertools import combinations

import numpy as np
import mcp.server.stdio
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.types import Tool, TextContent
from scipy import stats

# 自定义JSON编码器，处理numpy数据类型
class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(NumpyEncoder, self).default(obj)

def safe_json_dumps(obj, **kwargs):
    """安全的JSON序列化，处理numpy数据类型"""
    return json.dumps(obj, cls=NumpyEncoder, **kwargs)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('small_sample_mcp.log', encoding='utf-8'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

# 创建MCP服务器实例
server = Server("small-sample-nonparametric")

# ==================== 统计函数实现 ====================

def data_characteristics(data: np.ndarray) -> Dict[str, float]:
    """计算数据特征"""
    data = np.array(data)
    n = len(data)
    mean_val = float(np.mean(data))
    std_val = float(np.std(data, ddof=1))
    cv = std_val / abs(mean_val) if mean_val != 0 else float('inf')
    
    # 计算偏度和峰度
    centered = data - mean_val
    m2 = float(np.mean(centered**2))
    m3 = float(np.mean(centered**3))
    m4 = float(np.mean(centered**4))
    
    skewness = m3 / (m2**(3/2)) if m2 > 0 else 0.0
    kurtosis = m4 / (m2**2) - 3 if m2 > 0 else 0.0
    
    return {
        "sample_size": int(n),
        "mean": mean_val,
        "std": std_val,
        "cv": cv,
        "skewness": skewness,
        "kurtosis": kurtosis,
        "median": float(np.median(data)),
        "min": float(np.min(data)),
        "max": float(np.max(data))
    }



def wilcoxon_signed_rank_test(data: np.ndarray, median: float = 0) -> Dict[str, Any]:
    """Wilcoxon符号秩检验 (小样本精确)"""
    data = np.array(data)
    n = len(data)
    
    if n > 30:
        raise ValueError("样本量过大，请使用大样本近似方法")
    
    # 计算差值
    diff = data - median
    # 去除零值
    diff_nonzero = diff[diff != 0]
    n_nonzero = len(diff_nonzero)
    
    if n_nonzero == 0:
        return {
            "W_stat": 0,
            "p_value": 1.0,
            "median_estimate": np.median(data),
            "conf_interval": (np.median(data), np.median(data))
        }
    
    # 计算秩
    ranks = stats.rankdata(np.abs(diff_nonzero))
    
    # 计算正秩和
    W_plus = np.sum(ranks[diff_nonzero > 0])
    
    # 精确p值计算
    if n_nonzero <= 20:
        # 使用scipy.stats.wilcoxon进行精确计算
        try:
            # 使用scipy的wilcoxon函数，它会自动选择精确或近似方法
            _, p_value = stats.wilcoxon(diff_nonzero, alternative='two-sided')
        except:
            # 如果失败，使用正态近似
            mu = n_nonzero * (n_nonzero + 1) / 4
            sigma = sqrt(n_nonzero * (n_nonzero + 1) * (2 * n_nonzero + 1) / 24)
            z = (W_plus - mu) / sigma
            p_value = 2 * (1 - stats.norm.cdf(abs(z)))
    else:
        # 使用正态近似
        mu = n_nonzero * (n_nonzero + 1) / 4
        sigma = sqrt(n_nonzero * (n_nonzero + 1) * (2 * n_nonzero + 1) / 24)
        z = (W_plus - mu) / sigma
        p_value = 2 * (1 - stats.norm.cdf(abs(z)))
    
    # 点估计：Hodges-Lehmann估计量
    pairwise_means = []
    for i in range(n):
        for j in range(i, n):
            pairwise_means.append((data[i] + data[j]) / 2)
    median_estimate = np.median(pairwise_means)
    
    # 置信区间（简化版本）
    sorted_data = np.sort(data)
    if n >= 6:
        # 使用Wilcoxon符号秩检验的置信区间
        alpha = 0.05
        k = int(n * alpha / 2)
        conf_interval = (sorted_data[k], sorted_data[n-1-k])
    else:
        conf_interval = (sorted_data[0], sorted_data[-1])
    
    return {
        "W_stat": float(W_plus),
        "p_value": float(p_value),
        "median_estimate": float(median_estimate),
        "conf_interval": (float(conf_interval[0]), float(conf_interval[1]))
    }

def exact_generalized_signed_rank_test(data: np.ndarray, median: float = 0) -> Dict[str, Any]:
    """广义符号秩检验 (小样本精确)"""
    data = np.array(data)
    n = len(data)
    
    if n > 30:
        raise ValueError("样本量过大，请使用大样本近似方法")
    
    # 计算差值
    diff = data - median
    
    # 计算符号
    signs = np.sign(diff)
    positive_count = np.sum(signs > 0)
    negative_count = np.sum(signs < 0)
    zero_count = np.sum(signs == 0)
    
    # 有效样本量
    n_effective = positive_count + negative_count
    
    if n_effective == 0:
        return {
            "W_stat": 0,
            "p_value": 1.0,
            "median_estimate": np.median(data)
        }
    
    # 计算检验统计量（正号个数）
    W_stat = positive_count
    
    # 精确p值计算（二项分布）
    if n_effective <= 30:
        # 双侧检验
        p_value = 2 * min(stats.binom.cdf(W_stat, n_effective, 0.5),
                          1 - stats.binom.cdf(W_stat - 1, n_effective, 0.5))
        p_value = min(p_value, 1.0)
    else:
        # 正态近似
        mu = n_effective / 2
        sigma = sqrt(n_effective / 4)
        z = (W_stat - mu) / sigma
        p_value = 2 * (1 - stats.norm.cdf(abs(z)))
    
    return {
        "W_stat": int(W_stat),
        "p_value": float(p_value),
        "median_estimate": float(np.median(data)),
        "positive_count": int(positive_count),
        "negative_count": int(negative_count),
        "zero_count": int(zero_count)
    }

def exact_normal_scores_test(data: np.ndarray, median: float = 0) -> Dict[str, Any]:
    """正态记分检验 (小样本精确)"""
    data = np.array(data)
    n = len(data)
    
    if n > 30:
        raise ValueError("样本量过大，请使用大样本近似方法")
    
    # 计算差值
    diff = data - median
    diff_nonzero = diff[diff != 0]
    n_nonzero = len(diff_nonzero)
    
    if n_nonzero == 0:
        return {
            "W_stat": 0,
            "p_value": 1.0
        }
    
    # 计算秩
    ranks = stats.rankdata(np.abs(diff_nonzero))
    
    # 计算正态记分
    normal_scores = []
    for i in range(1, n_nonzero + 1):
        # 使用期望正态记分
        score = stats.norm.ppf(i / (n_nonzero + 1))
        normal_scores.append(score)
    
    # 计算检验统计量
    W_stat = 0
    for i, rank in enumerate(ranks):
        if diff_nonzero[i] > 0:
            W_stat += normal_scores[int(rank) - 1]
    
    # 计算方差（简化）
    var_W = sum(score**2 for score in normal_scores) / 4
    
    # 标准化统计量
    if var_W > 0:
        z_score = W_stat / sqrt(var_W)
        p_value = 2 * (1 - stats.norm.cdf(abs(z_score)))
    else:
        z_score = 0
        p_value = 1.0
    
    return {
        "W_stat": W_stat,
        "z_score": z_score,
        "p_value": p_value
    }

def exact_cox_stuart_test(data: np.ndarray) -> Dict[str, Any]:
    """Cox-Stuart趋势检验 (小样本精确)"""
    data = np.array(data)
    n = len(data)
    
    if n < 4:
        raise ValueError("样本量太小，至少需要4个观测值")
    
    # 将数据分成两半进行比较
    if n % 2 == 0:
        # 偶数个观测值
        first_half = data[:n//2]
        second_half = data[n//2:]
    else:
        # 奇数个观测值，去掉中间值
        first_half = data[:n//2]
        second_half = data[n//2+1:]
    
    # 计算配对差值的符号
    pairs = min(len(first_half), len(second_half))
    positive_signs = 0
    
    for i in range(pairs):
        if second_half[i] > first_half[i]:
            positive_signs += 1
    
    # 使用二项分布进行精确检验
    p_value = 2 * min(stats.binom.cdf(positive_signs, pairs, 0.5),
                      1 - stats.binom.cdf(positive_signs - 1, pairs, 0.5))
    p_value = min(p_value, 1.0)
    
    return {
        "pos": positive_signs,
        "n": pairs,
        "p_value": p_value
    }

def exact_runs_test(data: np.ndarray) -> Dict[str, Any]:
    """随机游程检验 (小样本精确)"""
    data = np.array(data)
    n = len(data)
    
    if n < 3:
        raise ValueError("样本量太小，至少需要3个观测值")
    
    # 计算中位数
    median = np.median(data)
    
    # 转换为二进制序列（大于中位数为1，小于为0，等于中位数的去除）
    binary_seq = []
    for x in data:
        if x > median:
            binary_seq.append(1)
        elif x < median:
            binary_seq.append(0)
    
    if len(binary_seq) < 3:
        return {
            "runs": 1,
            "p_value": 1.0
        }
    
    # 计算游程数
    runs = 1
    for i in range(1, len(binary_seq)):
        if binary_seq[i] != binary_seq[i-1]:
            runs += 1
    
    # 计算0和1的个数
    n1 = sum(binary_seq)
    n0 = len(binary_seq) - n1
    
    if n1 == 0 or n0 == 0:
        return {
            "runs": 1,
            "p_value": 1.0
        }
    
    # 精确p值计算（小样本）
    if len(binary_seq) <= 20:
        # 使用精确分布（简化版本）
        expected_runs = (2 * n0 * n1) / (n0 + n1) + 1
        variance_runs = (2 * n0 * n1 * (2 * n0 * n1 - n0 - n1)) / ((n0 + n1)**2 * (n0 + n1 - 1))
        
        if variance_runs > 0:
            z = (runs - expected_runs) / sqrt(variance_runs)
            p_value = 2 * (1 - stats.norm.cdf(abs(z)))
        else:
            p_value = 1.0
    else:
        # 正态近似
        expected_runs = (2 * n0 * n1) / (n0 + n1) + 1
        variance_runs = (2 * n0 * n1 * (2 * n0 * n1 - n0 - n1)) / ((n0 + n1)**2 * (n0 + n1 - 1))
        
        if variance_runs > 0:
            z = (runs - expected_runs) / sqrt(variance_runs)
            p_value = 2 * (1 - stats.norm.cdf(abs(z)))
        else:
            p_value = 1.0
    
    return {
        "runs": runs,
        "p_value": p_value,
        "n0": n0,
        "n1": n1
    }

# ==================== MCP服务器工具定义 ====================

@server.list_tools()
async def handle_list_tools() -> list[Tool]:
    """列出可用工具"""
    return [
        Tool(
            name="sample_size_check",
            description="检查样本大小是否适合小样本精确检验 (n <= 30)",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "样本数据"
                    }
                },
                "required": ["data"]
            }
        ),

        Tool(
            name="wilcoxon_signed_rank",
            description="Wilcoxon符号秩检验 (精确方法，含点估计和区间估计)",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "样本数据"
                    },
                    "median": {
                        "type": "number",
                        "description": "假设中位数",
                        "default": 0
                    }
                },
                "required": ["data"]
            }
        ),
        Tool(
            name="exact_generalized_signed_rank",
            description="广义符号秩检验 (小样本精确)",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "样本数据"
                    },
                    "median": {
                        "type": "number",
                        "description": "假设中位数",
                        "default": 0
                    }
                },
                "required": ["data"]
            }
        ),
        Tool(
            name="exact_normal_scores",
            description="正态记分检验 (小样本精确)",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "样本数据"
                    },
                    "median": {
                        "type": "number",
                        "description": "假设中位数",
                        "default": 0
                    }
                },
                "required": ["data"]
            }
        ),
        Tool(
            name="exact_cox_stuart_trend",
            description="Cox-Stuart趋势检验 (小样本精确)",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "时间序列数据"
                    }
                },
                "required": ["data"]
            }
        ),
        Tool(
            name="exact_runs_test",
            description="随机游程检验 (小样本精确)",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "样本数据"
                    }
                },
                "required": ["data"]
            }
        )
    ]

@server.call_tool()
async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> list[TextContent]:
    """处理工具调用"""
    try:
        if name == "sample_size_check":
            return await sample_size_check(arguments)

        elif name == "wilcoxon_signed_rank":
            return await wilcoxon_signed_rank(arguments)
        elif name == "exact_generalized_signed_rank":
            return await exact_generalized_signed_rank(arguments)
        elif name == "exact_normal_scores":
            return await exact_normal_scores(arguments)
        elif name == "exact_cox_stuart_trend":
            return await exact_cox_stuart_trend(arguments)
        elif name == "exact_runs_test":
            return await exact_runs_test_handler(arguments)
        else:
            return [TextContent(type="text", text=f"未知工具: {name}")]
    
    except Exception as e:
        logger.error(f"工具调用失败: {e}\n{traceback.format_exc()}")
        return [TextContent(type="text", text=f"工具调用失败: {str(e)}")]

# ==================== 工具处理函数 ====================

async def sample_size_check(args: Dict[str, Any]) -> list[TextContent]:
    """检查样本大小"""
    try:
        data = np.array(args["data"])
        n = len(data)
        
        if n <= 30:
            recommendation = "适合使用小样本精确检验"
            suitable = True
            methods = ["Wilcoxon符号秩检验", "广义符号秩检验", "正态记分检验"]
        else:
            recommendation = "样本量过大，建议使用大样本近似检验"
            suitable = False
            methods = []
            
        result = {
            "sample_size": n,
            "suitable_for_small_sample": suitable,
            "recommendation": recommendation,
            "threshold": 30,
            "available_methods": methods,
            "data_characteristics": data_characteristics(data)
        }
        
        return [TextContent(type="text", text=safe_json_dumps(result, ensure_ascii=False, indent=2))]
    
    except Exception as e:
        logger.error(f"样本大小检查失败: {e}")
        return [TextContent(type="text", text=f"样本大小检查失败: {str(e)}")]



async def wilcoxon_signed_rank(args: Dict[str, Any]) -> list[TextContent]:
    """Wilcoxon符号秩检验"""
    try:
        data = np.array(args["data"])
        median = args.get("median", 0)
        
        result = wilcoxon_signed_rank_test(data, median)
        
        response = {
            "test_name": "Wilcoxon符号秩检验 (小样本精确)",
            "null_hypothesis": f"中位数 = {median}",
            "sample_size": len(data),
            "test_statistic": result["W_stat"],
            "p_value": result["p_value"],
            "point_estimate": {
                "median_hodges_lehmann": result["median_estimate"]
            },
            "confidence_interval": {
                "lower": result["conf_interval"][0],
                "upper": result["conf_interval"][1],
                "confidence_level": 0.95
            },
            "conclusion": f"{'拒绝' if result['p_value'] < 0.05 else '不拒绝'}零假设 (α = 0.05)",
            "method_description": "基于符号秩的精确检验，适用于对称分布的位置参数检验"
        }
        
        return [TextContent(type="text", text=safe_json_dumps(response, ensure_ascii=False, indent=2))]

    except Exception as e:
        logger.error(f"Wilcoxon符号秩检验失败: {e}")
        return [TextContent(type="text", text=f"Wilcoxon符号秩检验失败: {str(e)}")]

async def exact_generalized_signed_rank(args: Dict[str, Any]) -> list[TextContent]:
    """广义符号秩检验 (精确)"""
    try:
        data = np.array(args["data"])
        median = args.get("median", 0)
        
        result = exact_generalized_signed_rank_test(data, median)
        
        response = {
            "test_name": "广义符号秩检验 (小样本精确)",
            "null_hypothesis": f"中位数 = {median}",
            "sample_size": len(data),
            "test_statistic": result["W_stat"],
            "p_value": result["p_value"],
            "sign_counts": {
                "positive": result["positive_count"],
                "negative": result["negative_count"],
                "zero": result["zero_count"]
            },
            "point_estimate": {
                "median": result["median_estimate"]
            },
            "conclusion": f"{'拒绝' if result['p_value'] < 0.05 else '不拒绝'}零假设 (α = 0.05)",
            "method_description": "基于符号的精确检验，对分布形状无要求，适用于任意连续分布"
        }
        
        return [TextContent(type="text", text=safe_json_dumps(response, ensure_ascii=False, indent=2))]

    except Exception as e:
        logger.error(f"广义符号秩检验失败: {e}")
        return [TextContent(type="text", text=f"广义符号秩检验失败: {str(e)}")]

async def exact_normal_scores(args: Dict[str, Any]) -> list[TextContent]:
    """正态记分检验 (精确)"""
    try:
        data = np.array(args["data"])
        median = args.get("median", 0)
        
        result = exact_normal_scores_test(data, median)
        
        response = {
            "test_name": "正态记分检验 (小样本精确)",
            "null_hypothesis": f"中位数 = {median}",
            "sample_size": len(data),
            "test_statistic": result["W_stat"],
            "z_score": result.get("z_score", "未计算"),
            "p_value": result["p_value"],
            "conclusion": f"{'拒绝' if result['p_value'] < 0.05 else '不拒绝'}零假设 (α = 0.05)",
            "method_description": "基于正态记分的精确检验，当数据近似正态分布时效果最佳"
        }
        
        return [TextContent(type="text", text=safe_json_dumps(response, ensure_ascii=False, indent=2))]

    except Exception as e:
        logger.error(f"正态记分检验失败: {e}")
        return [TextContent(type="text", text=f"正态记分检验失败: {str(e)}")]

async def exact_cox_stuart_trend(args: Dict[str, Any]) -> list[TextContent]:
    """Cox-Stuart趋势检验 (精确)"""
    try:
        data = np.array(args["data"])
        
        result = exact_cox_stuart_test(data)
        
        response = {
            "test_name": "Cox-Stuart趋势检验 (小样本精确)",
            "null_hypothesis": "数据无趋势",
            "sample_size": len(data),
            "positive_signs": result["pos"],
            "total_pairs": result["n"],
            "p_value": result["p_value"],
            "conclusion": f"{'存在' if result['p_value'] < 0.05 else '不存在'}显著趋势 (α = 0.05)",
            "method_description": "基于二项分布的精确趋势检验，检验时间序列是否存在单调趋势"
        }
        
        return [TextContent(type="text", text=safe_json_dumps(response, ensure_ascii=False, indent=2))]

    except Exception as e:
        logger.error(f"Cox-Stuart趋势检验失败: {e}")
        return [TextContent(type="text", text=f"Cox-Stuart趋势检验失败: {str(e)}")]

async def exact_runs_test_handler(args: Dict[str, Any]) -> list[TextContent]:
    """随机游程检验 (精确)"""
    try:
        data = np.array(args["data"])
        
        result = exact_runs_test(data)
        
        response = {
            "test_name": "随机游程检验 (小样本精确)",
            "null_hypothesis": "数据序列是随机的",
            "sample_size": len(data),
            "runs_count": result["runs"],
            "above_median": result.get("n1", "未计算"),
            "below_median": result.get("n0", "未计算"),
            "p_value": result["p_value"],
            "conclusion": f"数据序列{'不是' if result['p_value'] < 0.05 else '是'}随机的 (α = 0.05)",
            "method_description": "检验数据序列的随机性，基于游程数的精确分布"
        }
        
        return [TextContent(type="text", text=safe_json_dumps(response, ensure_ascii=False, indent=2))]

    except Exception as e:
        logger.error(f"随机游程检验失败: {e}")
        return [TextContent(type="text", text=f"随机游程检验失败: {str(e)}")]

# ==================== 主函数 ====================

async def main():
    """主函数：启动MCP服务器"""
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="small-sample-nonparametric",
                server_version="1.0.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )

if __name__ == "__main__":
    asyncio.run(main())